# -*- coding: utf-8 -*-

# base pkgs
import sys

# thirdparty pkgs
import numpy as np
from scipy.interpolate import UnivariateSpline

# private pkgs
from utils.crdlog import log
from algorithm.physical_property.tab_loader import TabLoader
from utils.numeric import gaussian_smooth
from utils.constant import MeshMode, Physics, BoundaryConditionType, InitVar

def cal_init_mesh(nx: int, value: int, mode: int) -> np.ndarray:
    """ 构建网格:

    Args:
        nx: int; 网格长度
        value: float; 初值
        mode: str; 1.'uniform'; 均匀网格
                  2. ... 待加入

    Returns:
        numpy.ndarray 网格
    """
    if mode == MeshMode.uniform.value:
        data = np.zeros(nx) + value
    return data

def cal_tilt_angle(h: np.ndarray, X: np.ndarray, is_smooth: bool = False) -> np.ndarray:
    """ 计算平滑后的斜率角度

    Args:
        h (np.ndarray): 每个位置点的高程
        X (np.ndarray): 每个位置点的里程

    Returns:
        np.ndarray: 平滑后的斜率角度
    """
    # 首先计算原始角度
    theta_segment = np.arctan2(np.diff(h), np.diff(X))
    if is_smooth:
        theta_segment_smooth = gaussian_smooth(theta_segment, 61)
        return theta_segment_smooth
    else:
        return theta_segment

def cal_cross_section_area(D_segment: np.ndarray) -> np.ndarray:
    """ 计算管道的横截面积

    Args:
        D_segment : np.ndarray  管道直径

    Returns:
        np.ndarray  管道的横截面积
    """
    return np.pi * (D_segment / 2) ** 2

def cal_boundary_velocity(physical_property_data: dict, inlet_S: float, A: float) -> tuple[float, float]:
    """由输入计算边界流速

    Args:
        physical_property_data (dict): 物性表
        inlet_S (float): 边界处质量流量
        A (float): 截面积

    Returns:
        tuple[float, float]: (inlet_V_g, inlet_V_L)
    """
    # # 稳态 pvt 中的全拼命名
    inlet_Rs_g = physical_property_data["gas_mass_fraction"][0]
    inlet_rho_L = physical_property_data["liquid_density"][0]
    inlet_rho_g = physical_property_data["gas_density"][0]

    # 入口速度
    inlet_V_L = (inlet_S * (1 - inlet_Rs_g) / inlet_rho_L + inlet_S * inlet_Rs_g / inlet_rho_g / Physics.R_D.value) / A
    inlet_V_g = inlet_V_L * Physics.R_D.value
    return inlet_V_g, inlet_V_L

def cal_boundary_alpha(physical_property_data: dict, inlet_V_g:float, inlet_V_L:float, inlet_S:float, A:float) -> tuple[float, float]:
    """计算气液相截面分数

    Args:
        physical_property_data(dict): 物性参数表
        inlet_V_g(float): 气相速度
        inlet_V_L(float): 液相速度
        inlet_S(float): 边界处质量流量
        A(float): 截面积

    Returns:
        tuple[float, float]: (inlet_alpha_g, inlet_alpha_L)
    """
    # 稳态 pvt 中的全拼命名
    inlet_Rs_g = physical_property_data["gas_mass_fraction"][0]
    inlet_rho_L = physical_property_data["liquid_density"][0]
    inlet_rho_g = physical_property_data["gas_density"][0]

    inlet_alpha_L = inlet_S * (1 - inlet_Rs_g) / inlet_rho_L / inlet_V_L / A
    inlet_alpha_g = inlet_S * inlet_Rs_g / inlet_rho_g / inlet_V_g / A
    return inlet_alpha_g, inlet_alpha_L

def judge_boundary_condition_class(params: dict) -> int:
    """ 判断边界条件类型, 目前支持三种边界条件的判定

    Args:
        params (dict): 存储输入边界条件的字典

    Returns:
        int: class id   
    """

    condition_inlet_S_outlet_p = [
        params["inlet_T_value"][0] != "None",
        params["inlet_M_value"][0] != "None",
        params["outlet_p_value"][0] != "None",
    ]
    condition_inlet_p_outlet_S = [
        params["inlet_p_value"][0] != "None",
        params["outlet_M_value"][0] != "None",
        params["outlet_T_value"][0] != "None",
    ]
    condition_inlet_p_outlet_p = [
        params["inlet_p_value"][0] != "None",
        params["outlet_p_value"][0] != "None",
        params["outlet_T_value"][0] != "None",
    ]
    
    if all(condition_inlet_S_outlet_p):
        return BoundaryConditionType.inlet_M_outlet_p.value
    elif all(condition_inlet_p_outlet_S):
        return BoundaryConditionType.inlet_p_outlet_M.value
    elif all(condition_inlet_p_outlet_p):
        return BoundaryConditionType.inlet_p_outlet_p.value
    else:
        raise log.error("please check the input, boundary condition is not correct")

def cal_and_save_boundary_condition(params: dict, physical_property_data: dict) -> dict:
    """ 根据工业边界条件计算CFD边界条件并存储到字典

    根据入口温度和出口压力计算入口边界条件,包括气液两相的密度、速度和截面分数,
    并将计算结果存储到参数字典中。
    ♥ 实际不管工业中的边界条件如何, 在CFD仿真计算中都需要转换成边界处的气液两相的密度、速度和截面分数;
    最终用于实际迭代中的束缚条件使方程有解.

    Args:
        params (dict): 存储计算参数的字典,包含入口温度、出口压力等边界条件
        physical_property_data (dict): 物性参数数据

    Returns:
        dict: 更新后的参数字典,包含计算得到的入口边界条件
    """
    if params["boundary_class"] == BoundaryConditionType.inlet_M_outlet_p.value:
        inlet_V_g, inlet_V_L = cal_boundary_velocity(physical_property_data, params["inlet_M_value"][0], (params["D_segment"][0] / 2) ** 2 * np.pi)
        inlet_alpha_g, inlet_alpha_L = cal_boundary_alpha(physical_property_data, inlet_V_g, inlet_V_L, params["inlet_M_value"][0], (params["D_segment"][0] / 2) ** 2 * np.pi)

        params["inlet_V_g"] = inlet_V_g
        params["inlet_V_L"] = inlet_V_L
        params["inlet_alpha_g"] = inlet_alpha_g
        params["inlet_alpha_L"] = inlet_alpha_L


    elif params["boundary_class"] == BoundaryConditionType.inlet_p_outlet_M.value:
        # 出口定流实际仍然是向内流, 速度应该取负号, 出口计算模型与入口一致
        outlet_V_g, outlet_V_L = - cal_boundary_velocity(physical_property_data, params["outlet_M_value"][0], (params["D_segment"][-1] / 2) ** 2 * np.pi) # 负号表示出口向里流
        outlet_alpha_g, outlet_alpha_L = cal_boundary_alpha(physical_property_data, outlet_V_g, outlet_V_L, params["outlet_M_value"][0], (params["D_segment"][-1] / 2) ** 2 * np.pi)

        params["outlet_V_g"] = outlet_V_g
        params["outlet_V_L"] = outlet_V_L
        params["outlet_alpha_g"] = outlet_alpha_g
        params["outlet_alpha_L"] = outlet_alpha_L

    elif params["boundary_class"] == BoundaryConditionType.inlet_p_outlet_p.value:
        outlet_rho_L = physical_property_data["liquid_density"][-1]
        outlet_rho_g = physical_property_data["gas_density"][-1]

        # 假定初始出口速度1m/s, 这个值假定在工业合理范围内即可
        outlet_V_g = 1
        outlet_V_L = 1
        params["outlet_M_value"] = (outlet_rho_g * outlet_V_g + outlet_rho_L * outlet_V_L) * (params["D_segment"][-1] / 2) ** 2 * np.pi
        outlet_alpha_g, outlet_alpha_L = cal_boundary_alpha(physical_property_data, outlet_V_g, outlet_V_L, params["outlet_M_value"][0], (params["D_segment"][-1] / 2) ** 2 * np.pi)

        params["outlet_V_g"] = outlet_V_g
        params["outlet_V_L"] = outlet_V_L
        params["outlet_alpha_g"] = outlet_alpha_g
        params["outlet_alpha_L"] = outlet_alpha_L

def _transfer_pipelength_X_h_to_codinate(initdata):
    # 转换为坐标形式

    # 高程
    h = initdata["pipe"]["h"]
    if isinstance(h, list):
        h = np.array(h, np.float64)
    else:
        h = np.asarray([h])

    # 里程 / 管长
    pipe_X_type = [ "X" in initdata["pipe"],
                    "L" in initdata["pipe"]]

    if pipe_X_type == [False, True]:
        # 仅输入管长
        L = initdata["pipe"]["L"]
        if isinstance(L, list):
            L_node = np.array(L, np.float64)
            if L_node[0] == 0:
                print('管长有误，不应该出现初始段为0的情况')
                raise
        else:
            L_node = np.asarray([L])
        X = (L_node ** 2 - h ** 2) ** 0.5
        X_node = np.zeros(len(X) + 1, np.float64)
        for i in range(0, len(X)):
            X_node[i + 1] = X[i] + X_node[i]
        h_node = np.zeros(len(h)+1, np.float64)
        for i in range(0, len(h)):
            h_node[i + 1] = h[i] + h_node[i]

    elif pipe_X_type == [True, False]:
        # 仅输入里程
        X = np.array(initdata["pipe"]["X"], np.float64)
        if X[0] == 0:
            # 坐标形式
            X_node = X
            if h[0] == 0:
                h_node = h
            else:
                h_node = h - h[0] # 如果高程首位坐标不为0，需将坐标原点平移到首位
                # todo: 25-01-16 如果高程首位坐标不为0，需将坐标原点平移到首位，但是否会影响海拔？？
        else:
            # 分段形式 转换为坐标形式
            X_node = np.zeros(len(X)+1, np.float64)
            for i in range(0, len(X)):
                X_node[i+1] = X[i] + X_node[i]
            h_node = np.zeros(len(h) + 1, np.float64)
            for i in range(0, len(h)):
                h_node[i + 1] = h[i] + h_node[i]
    else:
        # 其他情况不符合
        raise log.error('please check the input, pipe length/mileage is not correct')
    return X_node, h_node

def _calculate_segment_number(params, X_node, h_node, initdata, rate = None):
    """ 计算网格段数

    Args:
        params (dict): 存储计算参数的字典
        X_node (np.ndarray): 管道里程
        h_node (np.ndarray): 管道高程
        initdata (dict): 存储初始化数据的字典
        rate (float): 网格段数与管段数比率

    Returns:
        dict: 更新后的参数字典
    """

    # 考虑分段情况 dx
    X_seg, h_seg = np.diff(X_node), np.diff(h_node)
    L_seg = (X_seg**2 + h_seg**2)**0.5
    dx = initdata["cfd"]["dx"]

    # dx大于当前段长时，需要处理，否则网格跨段，引入较大误差
    # todo: 25 0109 此处定义的dx为每一段管段至少有1/rate个网格，可能会导致运行变慢
    if not rate:
        rate = 1 # 0.5 #0.1
    if dx > np.min(L_seg) * rate:
        dx = np.min(L_seg) * rate
    params["dx"] = dx

    segment_number = 0
    for i in range(1, len(X_node)):
        segment_number_m = int(round(
            np.sqrt((X_node[i] - X_node[i - 1]) ** 2 + (h_node[i] - h_node[i - 1]) ** 2)
            / params["dx"], 0
        ))

        segment_number = segment_number + segment_number_m
        params[f"node_number_{i}"] = segment_number_m + 1

    params["segment_number"] = segment_number
    params["node_number"] = segment_number + 1

    # 如果实际网格数量过少，则逐步增加网格数量(2倍数/次)
    if params["segment_number"] < 10:
        rate = rate / 2
        params = _calculate_segment_number(X_node, h_node, initdata, rate)

    return params

def _generate_mesh_coordinates(params, initdata):
    """生成网格坐标"""

    # 规范管长，里程和高程的输入
    X_node, h_node = _transfer_pipelength_X_h_to_codinate(initdata)
    params = _calculate_segment_number(params, X_node, h_node, initdata, rate = None)

    X_node_segment = []
    h_node_segment = []
    for i in range(1, len(X_node)):
        X_node_med = np.linspace(X_node[i - 1], X_node[i], params[f"node_number_{i}"])
        h_node_med = np.linspace(h_node[i - 1], h_node[i], params[f"node_number_{i}"])

        if i == 1:
            X_node_segment.append(X_node_med)
            h_node_segment.append(h_node_med)
        else:
            X_node_segment[0] = np.concatenate((X_node_segment[0], X_node_med[1:]))
            h_node_segment[0] = np.concatenate((h_node_segment[0], h_node_med[1:]))

    params["X_node"] = np.array(X_node_segment[0])
    params["h_node"] = np.array(h_node_segment[0])


def input_init(initdata) -> dict:
    """从配置文件读取输入参数，用于稳态计算

    Args:
        initdata: 输入参数字典
        pp_computer: 物性计算器

    Returns:
        dict: 存储转化后用于计算的初场、边界条件和数值模拟参数
    """
    params = {}

    # ----------------------------------------------------------------
    # 管道物性，几何，空间离散化以及方程组迭代控制等的存储
    # ----------------------------------------------------------------
    # 存储网格大小
    params["dx"] = initdata["cfd"]["dx"]

    # 生成网格坐标
    _generate_mesh_coordinates(params, initdata)

    # 管道直径
    params["D_segment"] = cal_init_mesh(params["segment_number"], initdata["pipe"]["D"], MeshMode.uniform.value)

    # 计算角度和截面积
    if "smooth" in initdata["pipe"]:
        params["tilt_angle_segment"] = cal_tilt_angle(params["h_node"], params["X_node"], is_smooth=initdata["pipe"]["smooth"])
    else:
        params["tilt_angle_segment"] = cal_tilt_angle(params["h_node"], params["X_node"])
    params["cross_section_area_segment"] = cal_cross_section_area(params["D_segment"])

    # 管道内表面粗糙度
    params["wall_roughness_segment"] = cal_init_mesh(params["segment_number"], initdata["pipe"]["eps"], MeshMode.uniform.value)

    # 管道热传导系数
    params["K"] = initdata["pipe"]["K"]

    # 环境温度
    params["T0"] = initdata["environment"]["T0"]

    # ----------------------------------------------------------------
    #   存储边界条件, in表示入口, out表示出口, 有随时间变化的量需要判断
    #   其是否为列表
    # ----------------------------------------------------------------
    # 处理入口边界条件
    inlet_vars = ["p", "M", "T"]
    for var in inlet_vars:
        value = initdata["inlet"][var]["value"]
        if isinstance(value, list):
            params[f"inlet_{var}_value"] = value
            params[f"inlet_{var}_time"] = initdata["inlet"][var]["time"]
        else:
            params[f"inlet_{var}_value"] = [value]

    # 处理出口边界条件
    outlet_vars = ["p", "M", "T"]
    for var in outlet_vars:
        value = initdata["outlet"][var]["value"]
        if isinstance(value, list):
            params[f"outlet_{var}_value"] = value
            params[f"outlet_{var}_time"] = initdata["outlet"][var]["time"]
        else:
            params[f"outlet_{var}_value"] = [value]

    # 网格化压力和温度用于计算物性
    params["p_node"] = cal_init_mesh(params["segment_number"] + 1, params["inlet_p_value"][0], MeshMode.uniform.value) if params["inlet_p_value"][0] is not None else cal_init_mesh(params["segment_number"] + 1, params["outlet_p_value"][0], MeshMode.uniform.value)
    params["T_node"] = cal_init_mesh(params["segment_number"] + 1, params["inlet_T_value"][0], MeshMode.uniform.value) if params["inlet_T_value"][0] is not None else cal_init_mesh(params["segment_number"] + 1, params["outlet_T_value"][0], MeshMode.uniform.value)

    # 计算物性
    params['tab_path'] = initdata['TAB']['tab_path']

    # 判断是否存在物性参数器，如果不存在则调用
    physical_property_calculator = TabLoader(params['tab_path'])
    physical_property_data = physical_property_calculator.get_physical_properties_data(params["p_node"], params["T_node"])

    # -----------------------------------------------------------
    # 边界条件类型判定
    # -----------------------------------------------------------
    # 判断边界条件类型
    bc_class = judge_boundary_condition_class(params)

    # 存储边界条件类型
    params["boundary_class"] = bc_class

    # 边界条件初始化，即计算边界处其他的量用于求解初始化，这一步必须在流动初始化之前完成
    cal_and_save_boundary_condition(params, physical_property_data)

    # ----------------------------------------------------------------
    # 方程组要求解的流动参数初始化以及离散化到空间网格
    # ----------------------------------------------------------------
    # 初始化流动参数到空间网格
    init_params_to_mesh = {
        'alpha_g_node': ('inlet_alpha_g', 'outlet_alpha_g', params["segment_number"] + 1),
        'alpha_L_node': ('inlet_alpha_L', 'outlet_alpha_L', params["segment_number"] + 1),
        'V_g_segment': ('inlet_V_g', 'outlet_V_g', params["segment_number"]),
        'V_L_segment': ('inlet_V_L', 'outlet_V_L', params["segment_number"]),
    }

    # 遍历初始化各参数
    for param_name, (inlet_key, outlet_key, size) in init_params_to_mesh.items():
        if inlet_key in params:
            inlet_val = eval(f'params["{inlet_key}"]')
        else:
            inlet_val = None
        if outlet_key in params:
            outlet_val = eval(f'params["{outlet_key}"]')
        else:
            outlet_val = None
        try:
            init_val = inlet_val if inlet_val is not None else outlet_val
            params[param_name] = cal_init_mesh(size, init_val, MeshMode.uniform.value)
        except:
            raise log.error(f"初始化参数{param_name}失败")

    # 计算和初始化液滴截面分数 TODO: 25-01-22 需要修改
    params["V_D_segment"] = cal_init_mesh(params["segment_number"], 0, MeshMode.uniform.value)
    params["alpha_D_node"] = 1 - params["alpha_L_node"] - params["alpha_g_node"]
    
    # ----------------------------------------------------------------
    # 通过实例化的物性表初始化相关物性参数，存储到字典
    # ----------------------------------------------------------------
    params["rho_g_node"] = physical_property_data["gas_density"]
    params["rho_L_node"] = physical_property_data["liquid_density"]
    # params["c_g_node"], params["c_L_node"] = cal_sound_velocity_node(physical_property_data)
    params["mu_g_node"] = cal_init_mesh(params["segment_number"] + 1,
    physical_property_data["gas_viscosity"],MeshMode.uniform.value)
    params["mu_L_node"] = cal_init_mesh(params["segment_number"] + 1,
    physical_property_data["liquid_viscosity"],MeshMode.uniform.value)
    surface_tension_gas_liquid = physical_property_data['liquid_surface_tension']
    params["sigma_node"] = cal_init_mesh(params["segment_number"] + 1,
    surface_tension_gas_liquid, MeshMode.uniform.value)
    params["Cp_g_node"] = physical_property_data["gas_heat_capacity"]
    params["Cp_L_node"] = physical_property_data["liquid_heat_capacity"]
    params["liquid_surface_tension_node"] = physical_property_data["liquid_surface_tension"]
    params["liquid_viscosity_node"] = physical_property_data["liquid_viscosity"]
    params["gas_viscosity_node"] = physical_property_data["gas_viscosity"]
    params["gas_mass_fraction_node"] = physical_property_data["gas_mass_fraction"]
    params["rho_g_derivate_p_node"] = physical_property_data["derivative_of_gas_density_w.r.t_pressure"]
    params["rho_L_derivate_p_node"] = physical_property_data["derivative_of_liquid_density_w.r.t_pressure"]
    params["rho_g_derivate_T_node"] = physical_property_data["derivative_of_gas_density_w.r.t_temperature"]
    params["rho_L_derivate_T_node"] = physical_property_data["derivative_of_liquid_density_w.r.t_temperature"]
    params["c_g_node"] = np.sqrt(1 / params["rho_g_derivate_p_node"])
    params["c_L_node"] = np.sqrt(1 / params["rho_L_derivate_p_node"])



    # -----------------------------------------------------------
    # 动态时间更新
    # -----------------------------------------------------------
    params["cfl"] = InitVar.cfl.value
    params["dt_min"] = initdata["cfd"]["dt_min"]
    params["dt_max"] = initdata["cfd"]["dt_max"]
    if params["inlet_V_L"] != "None":
        dt_cfl = params["cfl"] * params["dx"] / params["inlet_V_L"]
    else:
        dt_cfl = params["cfl"] * params["dx"] / params["outlet_V_L"]
        
    if params["dt_min"]:
        params['dt'] = np.min([params["dt_min"], dt_cfl])  # 取最dt
    else:
        params['dt'] = dt_cfl

    # 由dt和模拟时间决定时间步数
    params["max_time"] = initdata["cfd"]["max_time"]

    
    params["record_time"] = 0
    return params, physical_property_calculator

